# -*- coding:utf-8 -*-
from operator import attrgetter

import matplotlib.pyplot as plt

from numpy import genfromtxt  # genfromtxt函数创建数组表格数据
import numpy as np
from sklearn import datasets, linear_model

import QueryData as query;

# 建立回归模型
regr = linear_model.LinearRegression()

data = query.daily("002333.SZ");
data = sorted(data, key = lambda e:e.__getitem__('trade_date'), reverse=False)
# print(data)
x = []
y = []
currentIdx = 0
recordCount = range(len(data))
for i in recordCount:
  d = data[i]
  # x.append([d['low'],d['close'],d['change'], d['vol']])
  x.append([d['open'],d['high'],d['low'],d['close'],d['change'], d['vol']])
  y.append([data[i+1]['high']])
  currentIdx = i
  if currentIdx == 30:
    break

print("X:", x)
print("Y:", y)

regr.fit(x, y)
print("coefficients:", regr.coef_)  # b1,...,bp（与x相结合的各个参数）
print("intercept:", regr.intercept_)  # b0（截面）

dateData = []
y_real_high = []
y_real_low = []
y_pred_high = []
currentIdx = currentIdx +1
while currentIdx < len(data)-1:
  # 当天的数据，用于预测下一天
  d = data[currentIdx]
  # x_pred = [[d['low'],d['close'],d['change'], d['vol']]]
  x_pred = [[d['open'],d['high'],d['low'],d['close'],d['change'], d['vol']]]

  # 预测下一天
  y_pred = regr.predict(x_pred)
  y_pred_high.append(round(y_pred[0][0], 2))

  # 下一天的数据
  nd = data[currentIdx+1]
  y_real_high.append(nd['high'])
  y_real_low.append(nd['low'])
  dateData.append(nd['trade_date'])

  currentIdx = currentIdx + 1
  pass

print("date", dateData)
print("y_real_high", y_real_high)
print("y_pred_high", y_pred_high)
print("y_real_low", y_real_low)

plt.title('Analysis')
plt.plot(dateData, y_real_high, color='green', label='real accuracy')
plt.plot(dateData, y_pred_high, color='red', label='pre accuracy')
plt.plot(dateData, y_real_low, color='gray', label='lowest accuracy')
plt.legend()
plt.show()